package com.tencent.supersonic.chat.server.executor;

import java.sql.Connection;
import java.util.Collections;
import java.util.List;
import java.util.Map;

import org.apache.commons.lang3.StringUtils;
import org.mortbay.log.Log;
import org.springframework.util.CollectionUtils;

import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson2.JSON;
import com.tencent.supersonic.chat.api.pojo.response.QueryResult;
import com.tencent.supersonic.chat.server.agent.Agent;
import com.tencent.supersonic.chat.server.agent.AgentToolType;
import com.tencent.supersonic.chat.server.pojo.ExecuteContext;
import com.tencent.supersonic.chat.server.service.AgentService;
import com.tencent.supersonic.common.util.ContextUtils;
import com.tencent.supersonic.headless.api.pojo.request.DatabaseReq;
import com.tencent.supersonic.headless.api.pojo.response.DatabaseResp;
import com.tencent.supersonic.headless.api.pojo.response.QueryState;
import com.tencent.supersonic.headless.server.service.DatabaseService;

import ai.chat2db.spi.MetaData;
import ai.chat2db.spi.config.DBConfig;
import ai.chat2db.spi.sql.Chat2DBContext;
import ai.chat2db.spi.sql.ConnectInfo;
import ai.chat2db.spi.sql.IDriverManager;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.input.Prompt;
import dev.langchain4j.model.input.PromptTemplate;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.provider.ModelProvider;

public class PlainDDLExecutor implements ChatQueryExecutor {
	
	private static final String INSTRUCTION = ""
    		+ "### Instructions:\n"
    		+ "    Your task is to convert a question into a SQL query, given a %s database schema.\n"
    		+ "    Adhere to these rules:\n"
    		+ "    - **Deliberately go through the question and database schema word by word** to appropriately answer the question\n"
    		+ "    - **Use Table Aliases** to prevent ambiguity. For example, `SELECT table1.col1, table2.col1 FROM table1 JOIN table2 ON table1.id = table2.id`.\n"
    		+ "    - When creating a ratio, always cast the numerator as float\n"
    		+ "\n"
    		+ "### Input:\n"
    		+ "    Generate a SQL query that answers the question %s.\n"
    		+ "    This query will run on a database whose schema is represented in this string:\n"
    		+ "    %s"
    		+ "\n"
    		+ "-- BPAR_CODE_NUM_VAL.CODE_VAL can be joined with KPIS_REPT_KPI_EMP_2022.KPI_CCY and BPAR_CODE_NUM_VAL.par_code ='ZBBZ' \n"
    		+ "-- BPAR_CODE_NUM_VAL.CODE_VAL can be joined with KPIS_REPT_KPI_EMP_2022.KPI_INTRV and BPAR_CODE_NUM_VAL.par_code ='ZBSDBS' \n"
    		+ "\n"
    		+ "### Response:\n"
    		+ "    Based on your instructions, here is the SQL query I have generated to answer the question: %s:\n"
    		+ "    ```sql";

	@Override
	public QueryResult execute(ExecuteContext executeContext) {
		// TODO Auto-generated method stub
		if (!"PLAIN_TEXT_DDL".equals(executeContext.getParseInfo().getQueryMode())) {
            return null;
        }
		// String ddl = null;
		StringBuilder ddl = new StringBuilder();
		
		DatabaseService dbService = ContextUtils.getBean(DatabaseService.class);
		
		Agent agent = executeContext.getAgent();
		List<String> tools = agent.getTools(AgentToolType.DDL2SQL_LLM);
		String databaseId = null;
		String dbType = "mysql";
		if(!CollectionUtils.isEmpty(tools) && tools.size() == 1) {
			Map map = JSONObject.parseObject(tools.get(0), Map.class);
			databaseId = null != map.get("dataSetIds")? map.get("dataSetIds").toString() : "0";
		}
		if(!CollectionUtils.isEmpty(executeContext.getTables()) && 
				null!=databaseId) {
			ConnectInfo connectInfo = buildConnectionInfo(Long.valueOf(databaseId),dbService);
			MetaData meta = null;
			if(null!=connectInfo) {
				dbType = connectInfo.getDbType();
			    meta = Chat2DBContext.PLUGIN_MAP.get(connectInfo.getDbType()).getMetaData();
			}
			Connection connection = null;
			if(null!=meta) {
				connection = buildConn(connectInfo,meta);
			}
			if(null!=connection) {
				
				for(String tablename : executeContext.getTables()) {
					String s = meta.tableDDL(connection, connectInfo.getDatabaseName(), connectInfo.getUser(), tablename);
					if(StringUtils.isNoneBlank(s)) {
						ddl.append(s).append("\n");
					}
				}
			}
			
		}
		
		String promptStr = String.format(INSTRUCTION, dbType,executeContext.getQueryText(),ddl.toString(),
                executeContext.getQueryText());
		
		
		Log.info("promtStr:{}",promptStr);
        Prompt prompt = PromptTemplate.from(promptStr).apply(Collections.EMPTY_MAP);
        
        AgentService agentService = ContextUtils.getBean(AgentService.class);
        Agent chatAgent = agentService.getAgent(executeContext.getAgent().getId());
        
        Log.info("agent prompt:{}",chatAgent.getPromptConfig().getPromptTemplate());

        ChatLanguageModel chatLanguageModel = ModelProvider.getChatModel(chatAgent.getModelConfig());
        Response<AiMessage> response = chatLanguageModel.generate(prompt.toUserMessage());

        
        
        
        QueryResult result = new QueryResult();
        Log.info("result:{}",JSON.toJSONString(result));
        result.setQueryState(QueryState.SUCCESS);
        result.setQueryMode(executeContext.getParseInfo().getQueryMode());
        
        result.setTextResult(response.content().text());
        
        String text = result.getTextResult();
        Log.info("resp:{}",JSON.toJSONString(response.content()));
        if (text.indexOf(";")> 0) {
        	text = text.split(";")[0];
        }
        result.setTextResult(text);
        Log.info("result:{}",JSON.toJSONString(result));

        return result;
	}
	
	private ConnectInfo buildConnectionInfo(Long id, DatabaseService databaseService) {
		DatabaseResp dataSource = databaseService.getDatabase(id);
		ConnectInfo connectInfo = null;
		if(null!=dataSource) {
			connectInfo = new ConnectInfo();
			//oracle用户名大写
			if(dataSource.getType().toLowerCase().equals("oracle")) {
				connectInfo.setUser(dataSource.getUsername().toUpperCase());
			}else {
				connectInfo.setUser(dataSource.getUsername());
			}
			connectInfo.setPassword(dataSource.getPassword());
		    connectInfo.setDbType(dataSource.getType().toUpperCase());
		    connectInfo.setUrl(dataSource.getUrl());
		    connectInfo.setDatabase(dataSource.getDatabase());
		    connectInfo.setSchemaName(dataSource.getSchema());
		    connectInfo.setDatabaseName(dataSource.getDatabase());
		    connectInfo.setConsoleOwn(false);
		    connectInfo.setUrl(dataSource.getUrl());
		    connectInfo.setPort(StringUtils.isNotBlank(dataSource.getPort()) ? Integer.parseInt(dataSource.getPort()) : null);
		    connectInfo.setHost(dataSource.getHost());
		    DBConfig dbConfig = Chat2DBContext.getDBConfig(dataSource.getType().toUpperCase());
			if(null!=dbConfig) {
				connectInfo.setDriver(dbConfig.getDefaultDriverConfig().getJdbcDriverClass());
				connectInfo.setDriverConfig(dbConfig.getDefaultDriverConfig());
			}
		}
		return connectInfo;
	}
	
	private Connection buildConn(ConnectInfo connectInfo, MetaData meta) {
		
		try {
			
			if(null!=meta) {
				Connection connection = IDriverManager.getConnection(connectInfo.getUrl(), connectInfo.getUser(), connectInfo.getPassword(),
	                    connectInfo.getDriverConfig(), connectInfo.getExtendMap());
				return connection;
			}
		}catch(Exception e) {
		}
		
		return null;
		
	}
	
	

}
